package org.lanyu.springainovel.common.util;

import org.lanyu.springainovel.chat.entity.KnowledgeDocument;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class DocumentSplitSaveUtil {

    private int maxToken = 500;

    private VectorStore vectorStore;

    public DocumentSplitSaveUtil(int maxToken, VectorStore vectorStore) {
        this.maxToken = maxToken;
        this.vectorStore = vectorStore;
    }

    /**
     * 切分文档并保存到向量数据�?     *
     *
     * @param document 文档
     * @param metadata 元数�?     * @return 是否成功
     */
    public boolean splitAndSave(KnowledgeDocument document, Map<String, Object> metadata) {
        System.out.println("开始文档切分和保存过程...");
        System.out.println("向量存储类型: " + vectorStore.getClass().getName());

        ensureIndexExists(vectorStore);
        org.springframework.ai.document.Document aiDocument = new org.springframework.ai.document.Document(
                document.getContent(), metadata);
        System.out.println("文档内容长度: " + document.getContent().length() + " 字符");
        return splitAndSaveWithTokenSize(aiDocument, maxToken);
    }

    /**
     * 按指定token大小切分文档并保�?     *
     *
     * @param aiDocument 文档
     * @param tokenSize  token大小
     * @return 是否成功
     */
    private boolean splitAndSaveWithTokenSize(org.springframework.ai.document.Document aiDocument, int tokenSize) {
        System.out.println("使用 " + tokenSize + " token大小进行文档切分");

        // 在切分前再次确保索引存在
        ensureIndexExists(vectorStore);

        TokenTextSplitter splitter = TokenTextSplitter.builder()
                .withChunkSize(tokenSize)
                .withMinChunkLengthToEmbed(0)
                .withMaxNumChunks(1000)
                .build();
        List<org.springframework.ai.document.Document> list = splitter.split(aiDocument);

        // 分批保存到向量数据库，每批最�?2�?
        int batchSize = 32;
        int totalChunks = list.size();
        System.out.println("文档总块�? " + totalChunks + "，开始分批处�?..");

        List<org.springframework.ai.document.Document> failedChunks = new ArrayList<>();

        for (int i = 0; i < totalChunks; i += batchSize) {
            int endIndex = Math.min(i + batchSize, totalChunks);
            List<org.springframework.ai.document.Document> batch = list.subList(i, endIndex);
            List<org.springframework.ai.document.Document> batchWithTokenSize = new ArrayList<>();

            try {
                // 为每个文档块添加当前使用的token大小到metadata
                for (org.springframework.ai.document.Document doc : batch) {
                    // 创建新的元数据，将token大小添加进去
                    Map<String, Object> docMetadata = new HashMap<>(doc.getMetadata());
                    docMetadata.put("tokenSize", tokenSize);

                    // 创建一个新的文档对象，因为Document类没有提供直接修改metadata的方�?
                    org.springframework.ai.document.Document newDoc = new org.springframework.ai.document.Document(
                            doc.getId(),
                            doc.getText(),
                            docMetadata);

                    batchWithTokenSize.add(newDoc);
                }

                vectorStore.add(batchWithTokenSize);
                System.out.println("成功添加�?" + (i / batchSize + 1) + " 批，包含 " + batchWithTokenSize.size() + " 个文档块");
            } catch (Exception e) {
                System.err.println("添加�?" + (i / batchSize + 1) + " 批文档块时出�? " + e.getMessage());
                System.err.println("错误类型: " + e.getClass().getName());

                // 特殊处理索引不存在的情况
                if (e.getMessage() != null && e.getMessage().contains("Index not found")) {
                    System.err.println("检测到索引不存在问题，尝试重新初始�?..");
                    try {
                        // 强制重新初始化索�?
                        ensureIndexExists(vectorStore);
                        // 重试一次当前批�?
                        vectorStore.add(batchWithTokenSize);
                        System.err.println("索引重新初始化后，批�?" + (i / batchSize + 1) + " 添加成功");
                        continue; // 跳过失败处理，继续下一个批�?
                    } catch (Exception retryException) {
                        System.err.println("索引重新初始化后仍然失败: " + retryException.getMessage());
                    }
                }

                e.printStackTrace(); // 用于调试的完整堆栈跟�
                 failedChunks.addAll(batch);
            }
        }

        // 处理失败的文档块
        if (!failedChunks.isEmpty()) {
            System.out.println("�?" + failedChunks.size() + " 个文档块保存失败，尝试减小token大小重新切分");

            // 减小token大小进行重新切分
            int newTokenSize = tokenSize / 2;
            if (newTokenSize < 100) {
                System.err.println("token大小已经很小 (" + newTokenSize + ")，无法继续切分，放弃处理");
                return false;
            }

            boolean allSuccess = true;
            for (org.springframework.ai.document.Document failedDoc : failedChunks) {
                boolean success = splitAndSaveWithTokenSize(failedDoc, newTokenSize);
                if (!success) {
                    allSuccess = false;
                }
            }

            return allSuccess;
        }

        return true;
    }

    /**
     * 索引检查方�?- 现在主要是验证和日志记录
     * 实际的索引创建由VectorStoreManager在创建时完成
     */
    public void ensureIndexExists(VectorStore vectorStore) {
        try {
            // 检查是否为 ElasticsearchVectorStore 实例
            if (vectorStore instanceof org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStore) {
                System.out.println("使用Elasticsearch向量存储，索引检查通过");

                // 尝试添加空列表来验证索引状�?
                try {
                    vectorStore.add(java.util.Collections.emptyList());
                    System.out.println("Elasticsearch 向量存储索引验证通过");
                } catch (Exception e) {
                    // 如果索引验证失败，记录警告但继续执行
                    // 因为VectorStoreManager应该在创建时已经处理了索引创�?
                    System.out.println("索引验证警告（可能不影响使用�? " + e.getMessage());
                }
            } else {
                System.out.println("使用非Elasticsearch向量存储: " + vectorStore.getClass().getName());
            }
        } catch (Exception e) {
            System.err.println("索引检查过程中出现非致命错�? " + e.getMessage());
            // 不抛出异常，让流程继续执�?
            }
        }

    }

